#!/bin/bash
set -e

echo "=========================================="
echo "Wan2.1-T2V-14B 推理服务启动"
echo "=========================================="

# 检查NPU设备
echo "检查NPU设备..."
if command -v npu-smi &> /dev/null; then
    npu-smi info || echo "警告: 无法获取NPU信息"
else
    echo "警告: npu-smi 命令不可用"
fi

# 设置CANN环境变量
echo "设置CANN环境..."
if [ -d "/usr/local/Ascend/ascend-toolkit" ]; then
    source /usr/local/Ascend/ascend-toolkit/set_env.sh
    echo "✓ CANN环境已设置"
else
    echo "警告: CANN安装目录不存在"
fi

# 设置PYTHONPATH
export PYTHONPATH=/workspace/Wan2.1:/app:$PYTHONPATH

# 验证Python环境
echo "验证Python环境..."
python3 --version

echo "验证torch_npu..."
python3 -c "import torch; import torch_npu; print(f'PyTorch版本: {torch.__version__}'); print(f'NPU设备数: {torch.npu.device_count()}')" || {
    echo "警告: torch_npu验证失败"
}

# 验证Wan2代码
echo "验证Wan2代码..."
if [ -f "/workspace/Wan2.1/generate.py" ]; then
    echo "✓ Wan2代码就绪: /workspace/Wan2.1/"
    ls -la /workspace/Wan2.1/generate.py
else
    echo "✗ 错误: Wan2代码未找到"
    exit 1
fi

# 验证torchrun命令
echo "验证torchrun..."
which torchrun && echo "✓ torchrun命令可用" || echo "警告: torchrun命令未找到"

# 打印配置信息
echo "=========================================="
echo "配置信息:"
echo "  模型路径: ${MODEL_PATH:-/data2/models/Wan2.1-T2V-14B}"
echo "  Wan2代码路径: /workspace/Wan2.1"
echo "  FastAPI目录: /app"
echo "  输出目录: ${OUTPUT_DIR:-/app/outputs}"
echo "  日志级别: ${LOG_LEVEL:-INFO}"
echo "  端口: ${PORT:-8000}"
echo "  PYTHONPATH: $PYTHONPATH"
echo "=========================================="

# 创建必要的目录
mkdir -p /app/outputs /app/temp /app/logs

# 启动应用
echo "启动Wan2.1-T2V API服务..."
exec python3 -m uvicorn api.main:app --host 0.0.0.0 --port ${PORT:-8000} --app-dir /app
